/**
 * @ Author: luoqi
 * @ Create Time: 2024-04-17 18:01
 * @ Modified by: luoqi
 * @ Modified time: 2025-07-11 00:51
 * @ Description:
 */

#include <bsp.h>
#include "iap.h"
#include "ymodem.h"
#include "cli.h"

#if IAP_USING_DBG
#define IAP_DBG(...)  kprintf(__VA_ARGS__)
#else 
#define IAP_DBG(...) 
#endif

static int _cmd_iap_hdl(int argc, char **argv);

static IapObj _iap;
static YModem _ymodem;

static uint8_t buf[1024 + 6];

static int _putc(uint8_t c)
{
    return _iap.send(&c, 1);
}

static int _getc(uint8_t *c, size_t timeout_ms)
{
    return _iap.recv(c, 1, timeout_ms);
}

typedef void (*Jump)(void);
int iap_jump(size_t addr)
{
    Jump jmp;
    size_t appstack = *(volatile size_t *)IAP_APP_ADDR;
    if((appstack & 0x2ffe0000) != 0x20000000) {
        QSH(" #! Stack pointer is not valid, appstack: 0x%08X !\r\n", appstack);
        return -1;
    } else {
        QSH("\r\n Jump to app, addr: 0x%08X, appstack: 0x%08X ...\r\n", addr, appstack);
        volatile size_t jmp_addr = *(volatile size_t *)(addr + 4);
        jmp = (Jump)jmp_addr;
        __set_PSP(*(volatile size_t *)IAP_APP_ADDR);
        __set_CONTROL(0);
        __set_MSP(*(volatile size_t *)IAP_APP_ADDR);
        jmp();
    }
    return 0;
}

static int _ymodem_start_cb(const uint8_t *buf, size_t len)
{
    IAP_DBG("\r\n Firmware Download start...\r\n");
    bsp_blink_set(50, 500);
    int offset = 0;
    for(size_t i = 0; i < len; i++) {
        if((buf[i] != 0) && (buf[i] != 0x1a)) {
            offset++;
        } else {
            break;
        }
    }
    int sz = 0;
    for(size_t i = offset + 1; i < len - offset; i++) {
        if(buf[i] >= 48 && buf[i] <= 57) {
            sz = sz * 10 + (buf[i] - '0');
        } else {
            break;
        }
    }
    IAP_DBG(" FileName: %s, FileSize: %d\r\n", buf, sz);
    IAP_DBG(" Erase flash, addr: %08X, size: %d\r\n", IAP_APP_ADDR, sz - 2000);
    int err = _iap.mem_erase(IAP_APP_ADDR, sz - 2000);
    if(err < 0) {
        IAP_DBG(" Erase flash failed, addr: %d, size: %d, ecode: %d\r\n", IAP_APP_ADDR, sz - 2000, err);
    }
    return 0;
}

static int _ymodem_data_cb(const uint8_t *buf, size_t len)
{
    bsp_blink_set(500, 50);
    static size_t offset = 0;
    int err = 0;
    err = _iap.mem_write(IAP_APP_ADDR + offset, buf, len);
    if(err < 0) {
        IAP_DBG(" Flash write failed, addr: %d, sz: %d, ecode: %d\r\n", IAP_APP_ADDR + offset, len, err);
    }
    offset += len;
    return err;
}

static int _ymodem_end_cb(const uint8_t *buf, size_t len)
{
    UNUSED(buf);
    UNUSED(len);
    bsp_blink_set(500, 500);
    cli_echo("blink stat finish");
    bsp_delay_ms(10);
    IAP_DBG(" Firmware download finished\r\n");
    return 0;
}

int  iap_init(IapSend send, IapRecv recv, IapMemWrite mem_write, IapMemErase mem_erase)
{
    _iap.send = send;
    _iap.recv = recv;
    _iap.mem_write = mem_write;
    _iap.mem_erase = mem_erase;
    ymodem_init(&_ymodem, buf, _putc, _getc, _ymodem_start_cb, _ymodem_data_cb, _ymodem_end_cb);

    cli_add("iap", _cmd_iap_hdl, "iap commands");
    return 0;
}

int iap_exec()
{
    int err = ymodem_download(&_ymodem, 30);
    if(err < 0) {
        IAP_DBG("\r\n #! ymodem download failed, ecode: %d\r\n", err);
        return -1;
    } else {
        iap_jump(IAP_APP_ADDR);
    }
    return 0;
}

int _cmd_iap_hdl(int argc, char **argv)
{
    if(argc == 1) {
        QSH(" -jmp\r\n");
    } else if(ISARG(argv[1], "jmp") && argc == 2) {
        iap_jump(IAP_APP_ADDR);
    }
    return 0;
}
